package com.tool4j.graph.factory;

import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.alibaba.cloud.ai.graph.node.LlmNode;
import com.alibaba.cloud.ai.graph.node.QuestionClassifierNode;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.tool4j.entity.TableMetadata;
import com.tool4j.entity.params.JavaCodeConfig;
import com.tool4j.node.SqlParserNode;
import com.tool4j.tools.GenerateCodeTool;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.core.ParameterizedTypeReference;

import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import static com.alibaba.cloud.ai.graph.action.AsyncEdgeAction.edge_async;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;

public class GenerateCodeGraphFactory {

    public static CompiledGraph build(ChatClient chatClient) throws Exception {
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> keyStrategyHashMap = new HashMap<>();
            keyStrategyHashMap.put("chat_id", new ReplaceStrategy());
            keyStrategyHashMap.put("session_id", new ReplaceStrategy());
            keyStrategyHashMap.put("input", new ReplaceStrategy());
            keyStrategyHashMap.put("intent_type", new ReplaceStrategy());
            keyStrategyHashMap.put("reply_sql", new ReplaceStrategy());
            keyStrategyHashMap.put("config", new ReplaceStrategy());
            keyStrategyHashMap.put("llm_response", new ReplaceStrategy());
            keyStrategyHashMap.put("tool_result", new ReplaceStrategy());
            return keyStrategyHashMap;
        };
        StateGraph graph = new StateGraph("generateCodeGraph", keyStrategyFactory);
        // 分类节点
        QuestionClassifierNode intentClassifier = QuestionClassifierNode.builder()
                .chatClient(chatClient)
                .inputTextKey("input")
                .categories(List.of("sql", "feature"))
                .classificationInstructions(List.of("判断用户输入的是SQL语句还是原始需求。如果是SQL语句，返回'sql'，否则返回'feature'"))
                .outputKey("intent_type")
                .build();
        graph.addNode("intent", node_async(intentClassifier));


        LlmNode generateSqlNode = LlmNode.builder()
                .chatClient(chatClient)
                .systemPromptTemplate("你是一个需求分析器，擅长根据用户的原始需求，生成对应的建表SQL语句：")
                .userPromptTemplateKey("input")
                .outputKey("reply_sql")
                .stream(true)
                .build();

        //sql解析节点
        NodeAction extractNode = new SqlParserNode(chatClient);

        NodeAction generateCodeNode = state -> {
            Map<String, Object> config = (HashMap) state.value("config").orElse(null);
            List<TableMetadata> llmResponse = (List<TableMetadata>) state.value("llm_response").orElseThrow();
            String toolResult = GenerateCodeTool.generateCode(llmResponse, config);
            return Map.of("tool_result", toolResult);
        };

        graph.addNode("generateSql", node_async(generateSqlNode));
        graph.addNode("extract", node_async(extractNode));
        graph.addNode("generateCode", node_async(generateCodeNode));

        graph.addEdge(StateGraph.START, "intent");
        graph.addConditionalEdges("intent", edge_async(state -> {
            String intentRaw = (String) state.value("intent_type").orElse("");
            String intent = intentRaw;
            try {
                // 去除 Markdown code block
                intentRaw = intentRaw.trim();
                if (intentRaw.startsWith("```json")) {
                    intentRaw = intentRaw.replaceFirst("^```json", "").trim();
                }
                if (intentRaw.startsWith("```")) {
                    intentRaw = intentRaw.replaceFirst("^```", "").trim();
                }
                if (intentRaw.endsWith("```")) {
                    intentRaw = intentRaw.replaceAll("```$", "").trim();
                }
                // 解析 JSON
                ObjectMapper mapper = new ObjectMapper();
                JsonNode node = mapper.readTree(intentRaw);
                if (node.has("category_name")) {
                    intent = node.get("category_name").asText();
                }
            } catch (Exception e) {
            }
            return intent.contains("sql") ? "extract" : "generateSql";
        }), Map.of("extract", "extract", "generateSql", "generateSql"));

        graph.addEdge("generateSql", "extract");
        graph.addEdge("extract", "generateCode");
        graph.addEdge("generateCode", StateGraph.END);
        return graph.compile();
    }

}
